package com.zk.schedule;

import com.zk.schedule.config.ConfUtils;
import com.zk.schedule.config.ZkConfig;
import com.zk.schedule.schedule.ScheduleDataManager;
import com.zk.schedule.schedule.ScheduleServer;
import com.zk.schedule.utils.ScheduleUtil;
import com.zk.schedule.zk.Zk;
import org.apache.commons.lang3.BooleanUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.scheduling.Trigger;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.support.ScheduledMethodRunnable;

import java.lang.reflect.Method;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.Timer;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 调度器核心管理
 */
public class ZkScheduleManager extends ThreadPoolTaskScheduler implements ApplicationContextAware {
    private static final long serialVersionUID = 1L;
    protected static final transient Logger LOGGER = LoggerFactory.getLogger(ZkScheduleManager.class);
    /**
     * 心跳间隔
     */
    private static final int HEART_BEAT_INTERVAL = 2000;

    private ScheduleDataManager scheduleDataManager;

    /**
     * 当前调度服务的信息
     */
    protected ScheduleServer currentScheduleServer;

    private ZkConfig config = ConfUtils.loadConfig();

    /**
     * 是否注册成功
     */
    private boolean isScheduleServerRegister = false;

    private ApplicationContext applicationcontext;

    private Map<String, Boolean> isOwnerMap = new ConcurrentHashMap<>();

    protected Timer hearBeatTimer;
    protected Lock initLock = new ReentrantLock();
    protected Lock registerLock = new ReentrantLock();

    private volatile String errorMessage = "No config Zookeeper connect information";
    private InitialThread initialThread;
    private DestroyThread destroyThread;
    private Zk zk;

    static {
        new Zk(ConfUtils.loadConfig().getCxnString());
    }

    public Zk getZk() {
        if (zk == null) {
            zk = new Zk(ConfUtils.loadConfig().getCxnString());
        }
        return zk;
    }

    public ZkScheduleManager() {
        this.currentScheduleServer = ScheduleServer.createScheduleServer();
    }

    /**
     * 初始化
     * @throws Exception
     */
    public void init() throws Exception {
        if (this.initialThread != null) {
            this.initialThread.stopThread();
        }
        this.initLock.lock();
        try {
            this.scheduleDataManager = null;
            OpAdapter.setScheduleManager(this);
            if (this.zk != null) {
                this.zk = null;
            }
            this.zk = getZk();
            this.errorMessage = "Zookeeper connecting ......" + ConfUtils.loadConfig().getCxnString();
            initialThread = new InitialThread(this);
            initialThread.setName("ScheduleManager-initialThread");
            initialThread.start();
        } finally {
            this.initLock.unlock();
        }
    }

    /**
     * 重新初始化.会重新加载配置文件，重新生成server
     * @throws Exception
     */
    public void reInit() throws Exception {
        if (config.isStart() || !config.isStopSchedule() || this.hearBeatTimer != null) {
            throw new Exception("调度器有任务处理，不能重新初始化");
        }
        LOGGER.info("reInit server start...");
        config = ConfUtils.refreshConfig();
        this.init();
    }

    /**
     * 在Zk状态正常后回调数据初始化
     *
     * @throws Exception
     */
    public void initialData() throws Exception {
        this.scheduleDataManager = new ScheduleDataManager();
        if (config.isStart()) {
            // 注册调度管理器
            this.scheduleDataManager.registerServer(this.currentScheduleServer);
            if (hearBeatTimer == null) {
                hearBeatTimer = new Timer("ScheduleManager-" + this.currentScheduleServer.getUuid() + "-HearBeat");
            }
            hearBeatTimer.schedule(new HeartBeatTimerTask(this), 2000, HEART_BEAT_INTERVAL);
            LOGGER.info("init server success!!!, serverId:{}", currentScheduleServer.getUuid());
        }
    }

    /**
     * 停server,并不能把已经执行的任务立即结束!<br>
     *     <p>stopSchedule=true;start=false;</p>
     *     <p>从zk上的server注册列表里remove当前server</p>
     *     <p>停心跳timer</p>
     */
    public void stopServer() {
        if(!currentScheduleServer.isRegister()) {
            throw new RuntimeException("调度器没有注册，不能停止");
        }
        if(destroyThread != null) {
            destroyThread.stopThread();
        }
        destroyThread = new DestroyThread(this);
        destroyThread.setName("ScheduleManager-destroyThread");
        destroyThread.start();
    }

    public void rewriteScheduleInfo() throws Exception {
        registerLock.lock();
        try {
            if (config.isStopSchedule()) {
                LOGGER.debug("外部命令终止调度,不再注册调度服务.{}", currentScheduleServer.getUuid());
                return;
            }
            // 先发送心跳信息
            if (StringUtils.isNotBlank(errorMessage)) {
                this.currentScheduleServer.setDealInfoDesc(errorMessage);
            }
            if (!this.scheduleDataManager.refreshServer(this.currentScheduleServer)) {
                // 更新信息失败，清除内存数据后重新注册
                this.clearMemInfo();
                this.scheduleDataManager.registerServer(this.currentScheduleServer);
            }
            isScheduleServerRegister = true;
        } finally {
            registerLock.unlock();
        }
    }

    /**
     * 清除内存中所有的已经取得的数据和任务队列,在心态更新失败，或者发现注册中心的调度信息被删除
     */
    public void clearMemInfo() {

    }

    /**
     * 根据当前调度服务器的信息，重新计算分配所有的调度任务
     * 任务的分配是需要加锁，避免数据分配错误。为了避免数据锁带来的负面作用，通过版本号来达到锁的目的
     * <p/>
     * 1、获取任务状态的版本号
     * 2、获取所有的服务器注册信息和任务队列信息
     * 3、清除已经超过心跳周期的服务器注册信息
     * 3、重新计算任务分配
     * 4、更新任务状态的版本号【乐观锁】
     * 5、更新任务队列的分配信息
     *
     * @throws Exception
     */
    public void assignTask() throws Exception {
        scheduleDataManager.clearExpireServer();
        List<String> serverList = scheduleDataManager.loadServerNames();
        if (!scheduleDataManager.isLeader(this.currentScheduleServer.getUuid(), serverList)) {
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("不是负责任务分配的Leader,直接返回:{}", this.currentScheduleServer.getUuid());
            }
            return;
        }
        // 设置初始化成功标准，避免在leader转换的时候，新增的线程组初始化失败
        scheduleDataManager.assignTask(this.currentScheduleServer.getUuid(), serverList);
    }

    /**
     * 定时向数据配置中心更新当前服务器的心跳信息。 如果发现本次更新的时间如果已经超过了，服务器死亡的心跳周期，则不能在向服务器更新信息。
     * 而应该当作新的服务器，进行重新注册。
     *
     * @throws Exception
     */
    public void refreshScheduleServer() throws Exception {
        try {
            rewriteScheduleInfo();
            // 如果任务信息没有初始化成功，不做任务相关的处理
            if (!isScheduleServerRegister) {
                return;
            }
            // 重新分配任务
            this.assignTask();
        } catch (Throwable e) {
            // 清除内存中所有的已经取得的数据和任务队列,避免心跳线程失败时候导致的数据重复
            this.clearMemInfo();
            if (e instanceof Exception) {
                throw (Exception) e;
            } else {
                throw new Exception(e.getMessage(), e);
            }
        }
    }

    private Runnable taskWrapper(final Runnable task) {
        return new Runnable() {
            public void run() {
                ScheduledMethodRunnable scheduledMethodRunnable = (ScheduledMethodRunnable) task;
                Method targetMethod = scheduledMethodRunnable.getMethod();
                String[] beanNames = applicationcontext.getBeanNamesForType(targetMethod.getDeclaringClass());
                if (null != beanNames && StringUtils.isNotEmpty(beanNames[0])) {
                    String name = ScheduleUtil.getTaskNameFormBean(beanNames[0], targetMethod.getName());
                    boolean isOwner = false;
                    try {
                        if (!isScheduleServerRegister) {
                            Thread.sleep(1000);
                        }
                        if (getZk().getClient().isConnected()) {
                            isOwner = scheduleDataManager.isOwner(name, currentScheduleServer.getUuid());
                            isOwnerMap.put(name, isOwner);
                        } else {
                            // 如果zk不可用，使用历史数据
                            isOwner = BooleanUtils.toBoolean(isOwnerMap.get(name));
                        }
                    } catch (Exception e) {
                        LOGGER.error("Check task owner error.", e);
                    }
                    if (isOwner) {
                        task.run();
                        LOGGER.info("Cron job has been executed.");
                    }
                }
            }
        };
    }

    public ScheduleDataManager getScheduleDataManager() {
        return scheduleDataManager;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationcontext) throws BeansException {
        this.applicationcontext = applicationcontext;
    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, long period) {
        return super.scheduleAtFixedRate(taskWrapper(task), period);
    }

    public ScheduledFuture<?> schedule(Runnable task, Trigger trigger) {
        return super.schedule(taskWrapper(task), trigger);
    }

    public ScheduledFuture<?> schedule(Runnable task, Date startTime) {
        return super.schedule(taskWrapper(task), startTime);
    }

    public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, Date startTime, long period) {
        return super.scheduleAtFixedRate(taskWrapper(task), startTime, period);
    }

    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, Date startTime, long delay) {
        return super.scheduleWithFixedDelay(taskWrapper(task), startTime, delay);
    }

    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, long delay) {
        return super.scheduleWithFixedDelay(taskWrapper(task), delay);
    }

    public String getScheduleServerUuid() {
        return currentScheduleServer == null ? null : currentScheduleServer.getUuid();
    }

    public Map<String, Boolean> getIsOwnerMap() {
        return isOwnerMap;
    }

}

class HeartBeatTimerTask extends java.util.TimerTask {
    private static final Logger LOGGER = LoggerFactory.getLogger(HeartBeatTimerTask.class);
    private ZkScheduleManager manager;

    public HeartBeatTimerTask(ZkScheduleManager scheduleManager) {
        manager = scheduleManager;
    }

    public void run() {
        try {
            Thread.currentThread().setPriority(Thread.MAX_PRIORITY);
            manager.refreshScheduleServer();
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
    }
}

class InitialThread extends Thread {
    private static final Logger LOGGER = LoggerFactory.getLogger(InitialThread.class);
    private ZkScheduleManager sm;

    public InitialThread(ZkScheduleManager sm) {
        this.sm = sm;
    }

    private boolean isStop = false;

    public void stopThread() {
        this.isStop = true;
    }

    @Override
    public void run() {
        if (isStop) {
            return;
        }
        sm.initLock.lock();
        try {
            sm.initialData();
        } catch (Throwable e) {
            LOGGER.error(e.getMessage(), e);
        } finally {
            sm.initLock.unlock();
        }
    }

}

class DestroyThread extends Thread {
    private static final Logger LOGGER = LoggerFactory.getLogger(DestroyThread.class);
    private ZkScheduleManager sm;

    private boolean isStop = false;

    public void stopThread() {
        this.isStop = true;
    }

    public DestroyThread(ZkScheduleManager sm) {
        this.sm = sm;
    }

    @Override
    public void run() {
        if (isStop) {
            return;
        }
        LOGGER.info("stop server start... serverId:{}", sm.currentScheduleServer.getUuid());
        sm.registerLock.lock();
        ZkConfig config = ConfUtils.loadConfig();
        try {
            config.setStopSchedule(true);
            config.setStart(false);
            sm.getScheduleDataManager().removeServer(sm.currentScheduleServer);
            sm.hearBeatTimer.cancel(); // 如果server从zk上移除成功，停止心跳
            sm.hearBeatTimer = null;
            LOGGER.info("stop server success!!! serverId:{}", sm.currentScheduleServer.getUuid());
        } catch (Throwable t) {
            config.setStopSchedule(false);
            config.setStart(true);
            throw new RuntimeException(t.getMessage(), t);
        } finally {
            sm.registerLock.unlock();
        }
    }

}